/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Haystack Software Inc. All rights reserved.
 *  Licensed under the PolyForm Strict License 1.0.0. See License.txt in the project root for
 *  license information.
 *--------------------------------------------------------------------------------------------*/

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See code-license.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import { assertNever } from "vs/base/common/assert"
import { DeferredPromise } from "vs/base/common/async"
import { CancellationToken } from "vs/base/common/cancellation"
import { SetMap } from "vs/base/common/map"
import { onUnexpectedExternalError } from "vs/base/common/errors"
import { IDisposable } from "vs/base/common/lifecycle"
import { ISingleEditOperation } from "vs/editor/common/core/editOperation"
import { Position } from "vs/editor/common/core/position"
import { Range } from "vs/editor/common/core/range"
import { LanguageFeatureRegistry } from "vs/editor/common/languageFeatureRegistry"
import {
  Command,
  InlineCompletion,
  InlineCompletionContext,
  InlineCompletionProviderGroupId,
  InlineCompletions,
  InlineCompletionsProvider,
} from "vs/editor/common/languages"
import { ILanguageConfigurationService } from "vs/editor/common/languages/languageConfigurationRegistry"
import { ITextModel } from "vs/editor/common/model"
import { fixBracketsInLine } from "vs/editor/common/model/bracketPairsTextModelPart/fixBrackets"
import { SingleTextEdit } from "vs/editor/common/core/textEdit"
import { getReadonlyEmptyArray } from "vs/editor/contrib/inlineCompletions/browser/utils"
import {
  SnippetParser,
  Text,
} from "vs/editor/contrib/snippet/browser/snippetParser"

export async function provideInlineCompletions(
  registry: LanguageFeatureRegistry<InlineCompletionsProvider>,
  position: Position,
  model: ITextModel,
  context: InlineCompletionContext,
  token: CancellationToken = CancellationToken.None,
  languageConfigurationService?: ILanguageConfigurationService,
): Promise<InlineCompletionProviderResult> {
  // Important: Don't use position after the await calls, as the model could have been changed in the meantime!
  const defaultReplaceRange = getDefaultRange(position, model)
  const providers = registry.all(model)

  const multiMap = new SetMap<
    InlineCompletionProviderGroupId,
    InlineCompletionsProvider<any>
  >()
  for (const provider of providers) {
    if (provider.groupId) {
      multiMap.add(provider.groupId, provider)
    }
  }

  function getPreferredProviders(
    provider: InlineCompletionsProvider<any>,
  ): InlineCompletionsProvider<any>[] {
    if (!provider.yieldsToGroupIds) {
      return []
    }
    const result: InlineCompletionsProvider<any>[] = []
    for (const groupId of provider.yieldsToGroupIds || []) {
      const providers = multiMap.get(groupId)
      for (const p of providers) {
        result.push(p)
      }
    }
    return result
  }

  type Result = Promise<InlineCompletions<InlineCompletion> | null | undefined>
  const states = new Map<
    InlineCompletionsProvider<InlineCompletions<InlineCompletion>>,
    Result
  >()

  const seen = new Set<
    InlineCompletionsProvider<InlineCompletions<InlineCompletion>>
  >()
  function findPreferredProviderCircle(
    provider: InlineCompletionsProvider<any>,
    stack: InlineCompletionsProvider[],
  ): InlineCompletionsProvider[] | undefined {
    stack = [...stack, provider]
    if (seen.has(provider)) {
      return stack
    }

    seen.add(provider)
    try {
      const preferred = getPreferredProviders(provider)
      for (const p of preferred) {
        const c = findPreferredProviderCircle(p, stack)
        if (c) {
          return c
        }
      }
    } finally {
      seen.delete(provider)
    }
    return undefined
  }

  function processProvider(provider: InlineCompletionsProvider<any>): Result {
    const state = states.get(provider)
    if (state) {
      return state
    }

    const circle = findPreferredProviderCircle(provider, [])
    if (circle) {
      onUnexpectedExternalError(
        new Error(
          `Inline completions: cyclic yield-to dependency detected. Path: ${circle.map((s) => (s.toString ? s.toString() : "" + s)).join(" -> ")}`,
        ),
      )
    }

    const deferredPromise = new DeferredPromise<
      InlineCompletions<InlineCompletion> | null | undefined
    >()
    states.set(provider, deferredPromise.p)
    ;(async () => {
      if (!circle) {
        const preferred = getPreferredProviders(provider)
        for (const p of preferred) {
          const result = await processProvider(p)
          if (result && result.items.length > 0) {
            // Skip provider
            return undefined
          }
        }
      }

      try {
        const completions = await provider.provideInlineCompletions(
          model,
          position,
          context,
          token,
        )
        return completions
      } catch (e) {
        onUnexpectedExternalError(e)
        return undefined
      }
    })().then(
      (c) => deferredPromise.complete(c),
      (e) => deferredPromise.error(e),
    )

    return deferredPromise.p
  }

  const providerResults = await Promise.all(
    providers.map(async (provider) => ({
      provider,
      completions: await processProvider(provider),
    })),
  )

  const itemsByHash = new Map<string, InlineCompletionItem>()
  const lists: InlineCompletionList[] = []
  for (const result of providerResults) {
    const completions = result.completions
    if (!completions) {
      continue
    }
    const list = new InlineCompletionList(completions, result.provider)
    lists.push(list)

    for (const item of completions.items) {
      const inlineCompletionItem = InlineCompletionItem.from(
        item,
        list,
        defaultReplaceRange,
        model,
        languageConfigurationService,
      )
      itemsByHash.set(inlineCompletionItem.hash(), inlineCompletionItem)
    }
  }

  return new InlineCompletionProviderResult(
    Array.from(itemsByHash.values()),
    new Set(itemsByHash.keys()),
    lists,
  )
}

export class InlineCompletionProviderResult implements IDisposable {
  constructor(
    /**
     * Free of duplicates.
     */
    public readonly completions: readonly InlineCompletionItem[],
    private readonly hashs: Set<string>,
    private readonly providerResults: readonly InlineCompletionList[],
  ) {}

  public has(item: InlineCompletionItem): boolean {
    return this.hashs.has(item.hash())
  }

  dispose(): void {
    for (const result of this.providerResults) {
      result.removeRef()
    }
  }
}

/**
 * A ref counted pointer to the computed `InlineCompletions` and the `InlineCompletionsProvider` that
 * computed them.
 */
export class InlineCompletionList {
  private refCount = 1
  constructor(
    public readonly inlineCompletions: InlineCompletions,
    public readonly provider: InlineCompletionsProvider,
  ) {}

  addRef(): void {
    this.refCount++
  }

  removeRef(): void {
    this.refCount--
    if (this.refCount === 0) {
      this.provider.freeInlineCompletions(this.inlineCompletions)
    }
  }
}

export class InlineCompletionItem {
  public static from(
    inlineCompletion: InlineCompletion,
    source: InlineCompletionList,
    defaultReplaceRange: Range,
    textModel: ITextModel,
    languageConfigurationService: ILanguageConfigurationService | undefined,
  ) {
    let insertText: string
    let snippetInfo: SnippetInfo | undefined
    let range = inlineCompletion.range
      ? Range.lift(inlineCompletion.range)
      : defaultReplaceRange

    if (typeof inlineCompletion.insertText === "string") {
      insertText = inlineCompletion.insertText

      if (
        languageConfigurationService &&
        inlineCompletion.completeBracketPairs
      ) {
        insertText = closeBrackets(
          insertText,
          range.getStartPosition(),
          textModel,
          languageConfigurationService,
        )

        // Modify range depending on if brackets are added or removed
        const diff = insertText.length - inlineCompletion.insertText.length
        if (diff !== 0) {
          range = new Range(
            range.startLineNumber,
            range.startColumn,
            range.endLineNumber,
            range.endColumn + diff,
          )
        }
      }

      snippetInfo = undefined
    } else if ("snippet" in inlineCompletion.insertText) {
      const preBracketCompletionLength =
        inlineCompletion.insertText.snippet.length

      if (
        languageConfigurationService &&
        inlineCompletion.completeBracketPairs
      ) {
        inlineCompletion.insertText.snippet = closeBrackets(
          inlineCompletion.insertText.snippet,
          range.getStartPosition(),
          textModel,
          languageConfigurationService,
        )

        // Modify range depending on if brackets are added or removed
        const diff =
          inlineCompletion.insertText.snippet.length -
          preBracketCompletionLength
        if (diff !== 0) {
          range = new Range(
            range.startLineNumber,
            range.startColumn,
            range.endLineNumber,
            range.endColumn + diff,
          )
        }
      }

      const snippet = new SnippetParser().parse(
        inlineCompletion.insertText.snippet,
      )

      if (
        snippet.children.length === 1 &&
        snippet.children[0] instanceof Text
      ) {
        insertText = snippet.children[0].value
        snippetInfo = undefined
      } else {
        insertText = snippet.toString()
        snippetInfo = {
          snippet: inlineCompletion.insertText.snippet,
          range: range,
        }
      }
    } else {
      assertNever(inlineCompletion.insertText)
    }

    return new InlineCompletionItem(
      insertText,
      inlineCompletion.command,
      range,
      insertText,
      snippetInfo,
      inlineCompletion.additionalTextEdits || getReadonlyEmptyArray(),
      inlineCompletion,
      source,
    )
  }

  constructor(
    readonly filterText: string,
    readonly command: Command | undefined,
    readonly range: Range,
    readonly insertText: string,
    readonly snippetInfo: SnippetInfo | undefined,

    readonly additionalTextEdits: readonly ISingleEditOperation[],

    /**
     * A reference to the original inline completion this inline completion has been constructed from.
     * Used for event data to ensure referential equality.
     */
    readonly sourceInlineCompletion: InlineCompletion,

    /**
     * A reference to the original inline completion list this inline completion has been constructed from.
     * Used for event data to ensure referential equality.
     */
    readonly source: InlineCompletionList,
  ) {
    filterText = filterText.replace(/\r\n|\r/g, "\n")
    insertText = filterText.replace(/\r\n|\r/g, "\n")
  }

  public withRange(updatedRange: Range): InlineCompletionItem {
    return new InlineCompletionItem(
      this.filterText,
      this.command,
      updatedRange,
      this.insertText,
      this.snippetInfo,
      this.additionalTextEdits,
      this.sourceInlineCompletion,
      this.source,
    )
  }

  public hash(): string {
    return JSON.stringify({
      insertText: this.insertText,
      range: this.range.toString(),
    })
  }

  public toSingleTextEdit(): SingleTextEdit {
    return new SingleTextEdit(this.range, this.insertText)
  }
}

export interface SnippetInfo {
  snippet: string
  /* Could be different than the main range */
  range: Range
}

function getDefaultRange(position: Position, model: ITextModel): Range {
  const word = model.getWordAtPosition(position)
  const maxColumn = model.getLineMaxColumn(position.lineNumber)
  // By default, always replace up until the end of the current line.
  // This default might be subject to change!
  return word
    ? new Range(
        position.lineNumber,
        word.startColumn,
        position.lineNumber,
        maxColumn,
      )
    : Range.fromPositions(position, position.with(undefined, maxColumn))
}

function closeBrackets(
  text: string,
  position: Position,
  model: ITextModel,
  languageConfigurationService: ILanguageConfigurationService,
): string {
  const lineStart = model
    .getLineContent(position.lineNumber)
    .substring(0, position.column - 1)
  const newLine = lineStart + text

  const newTokens = model.tokenization.tokenizeLineWithEdit(
    position,
    newLine.length - (position.column - 1),
    text,
  )
  const slicedTokens = newTokens?.sliceAndInflate(
    position.column - 1,
    newLine.length,
    0,
  )
  if (!slicedTokens) {
    return text
  }

  const newText = fixBracketsInLine(slicedTokens, languageConfigurationService)

  return newText
}
